{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Additional tips and tricks for designing networks\n",
    "\n",
    "This tutorial assumes that you have read\n",
    "the `network_design` tutorial,\n",
    "and have designed a network or two.\n",
    "Here, we will give a few advanced tips and tricks\n",
    "for designing networks that can be reused flexibly.\n",
    "In particular, these tips will use the\n",
    "`config` system, so we will also assume that\n",
    "you have gone over the `config` tutorial.\n",
    "\n",
    "Briefly, the general principles covered\n",
    "in this tutorial are\n",
    "\n",
    "0. Accept `**kwargs` to pass through network arguments\n",
    "0. Accept a config argument for groups of parameters\n",
    "\n",
    "We will demonstrate these principles\n",
    "using the two examples from the `network_design` tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import nengo\n",
    "from nengo.dists import Choice\n",
    "from nengo.processes import Piecewise\n",
    "from nengo.utils.ipython import hide_input\n",
    "\n",
    "\n",
    "def test_integrators(net):\n",
    "    with net:\n",
    "        piecewise = Piecewise({\n",
    "            0: 0,\n",
    "            0.2: 0.5,\n",
    "            1: 0,\n",
    "            2: -1,\n",
    "            3: 0,\n",
    "            4: 1,\n",
    "            5: 0\n",
    "        })\n",
    "        piecewise_inp = nengo.Node(piecewise)\n",
    "        nengo.Connection(piecewise_inp, net.pre_integrator.input)\n",
    "        input_probe = nengo.Probe(piecewise_inp)\n",
    "        pre_probe = nengo.Probe(net.pre_integrator.ensemble, synapse=0.01)\n",
    "        post_probe = nengo.Probe(net.post_integrator.ensemble, synapse=0.01)\n",
    "    with nengo.Simulator(net) as sim:\n",
    "        sim.run(6)\n",
    "    plt.figure()\n",
    "    plt.plot(sim.trange(), sim.data[input_probe], color='k')\n",
    "    plt.plot(sim.trange(), sim.data[pre_probe], color='b')\n",
    "    plt.plot(sim.trange(), sim.data[post_probe], color='g')\n",
    "\n",
    "\n",
    "hide_input()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Accept a `**kwargs` argument\n",
    "\n",
    "The standard `nengo.Network` accepts a number of arguments,\n",
    "including the widely used `seed` and `label` arguments.\n",
    "Sometimes it is helpful to be able to\n",
    "set these on your custom networks too.\n",
    "While there is nothing wrong\n",
    "with explicitly passing these arguments along,\n",
    "it is less typing to use the Python `**kwargs` construct.\n",
    "This special argument allows a function\n",
    "to accept any number of keyword arguments\n",
    "which we can then pass into the `Network` constructor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Integrator(n_neurons, dimensions, tau=0.1, **kwargs):\n",
    "    with nengo.Network(**kwargs) as net:\n",
    "        net.input = nengo.Node(size_in=dimensions)\n",
    "        net.ensemble = nengo.Ensemble(n_neurons, dimensions=dimensions)\n",
    "        nengo.Connection(net.ensemble, net.ensemble, synapse=tau)\n",
    "        nengo.Connection(net.input, net.ensemble, synapse=None, transform=tau)\n",
    "    return net\n",
    "\n",
    "\n",
    "net = nengo.Network(label=\"Two integrators\")\n",
    "with net:\n",
    "    # Make both integrators use LIFRate neurons\n",
    "    net.config[nengo.Ensemble].neuron_type = nengo.LIFRate()\n",
    "    net.pre_integrator = Integrator(50, 1, label='pre')\n",
    "    net.post_integrator = Integrator(50, 1, label='post')\n",
    "    nengo.Connection(net.pre_integrator.ensemble, net.post_integrator.input)\n",
    "test_integrators(net)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('pre integrator label:', net.pre_integrator.label)\n",
    "print('post integrator label:', net.post_integrator.label)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Accept a config argument for groups of parameters\n",
    "\n",
    "Often, you will not want to use the\n",
    "network-level defaults for all of your objects.\n",
    "Some objects need certain things overwritten,\n",
    "while others need other values overwritten.\n",
    "Again, it is possible to deal with this issue\n",
    "by adding more and more parameters,\n",
    "but this quickly gets out of hand.\n",
    "Instead, add a small number of arguments\n",
    "that optionally accept a `config` object,\n",
    "which allows for setting multiple parameters at once.\n",
    "\n",
    "In the coupled integrator network example,\n",
    "we make two connections.\n",
    "We have to be careful changing the defaults\n",
    "for those connections, as they are wildly different;\n",
    "one is a recurrent connection from an ensemble to itself,\n",
    "while the other is a connection from a node to an ensemble.\n",
    "We will accept a `config` object for the recurrent connection\n",
    "to make this easier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ConfigurableIntegrator(\n",
    "        n_neurons, dimensions, recurrent_config=None, **kwargs):\n",
    "    net = nengo.Network(**kwargs)\n",
    "    if recurrent_config is None:\n",
    "        recurrent_config = nengo.Config(nengo.Connection)\n",
    "        recurrent_config[nengo.Connection].synapse = nengo.Lowpass(0.1)\n",
    "    with net:\n",
    "        net.input = nengo.Node(size_in=dimensions)\n",
    "        net.ensemble = nengo.Ensemble(n_neurons, dimensions=dimensions)\n",
    "        with recurrent_config:\n",
    "            nengo.Connection(net.ensemble, net.ensemble)\n",
    "            tau = nengo.Config.default(nengo.Connection, 'synapse').tau\n",
    "        nengo.Connection(net.input, net.ensemble, synapse=None, transform=tau)\n",
    "    return net\n",
    "\n",
    "\n",
    "net = nengo.Network(label=\"Two integrators\")\n",
    "with net:\n",
    "    # Make both integrators use LIFRate neurons\n",
    "    net.config[nengo.Ensemble].neuron_type = nengo.LIFRate()\n",
    "    net.pre_integrator = ConfigurableIntegrator(50, 1)\n",
    "    # Give the post_integrator a shorter tau (should make integration fail)\n",
    "    recurrent_config = nengo.Config(nengo.Connection)\n",
    "    recurrent_config[nengo.Connection].synapse = nengo.Lowpass(0.01)\n",
    "    net.post_integrator = ConfigurableIntegrator(\n",
    "        50, 1, recurrent_config=recurrent_config)\n",
    "    nengo.Connection(net.pre_integrator.ensemble, net.post_integrator.input)\n",
    "test_integrators(net)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Longer example: double integrator network\n",
    "\n",
    "Recall in the previous tutorial that\n",
    "we created a model\n",
    "that released a lever 0.6 to 1.0 seconds\n",
    "after pressing a lever.\n",
    "Let's use the above principles,\n",
    "and the `config` system in general,\n",
    "to improve the code constructing this model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def controlled_integrator(n_neurons,\n",
    "                          dimensions,\n",
    "                          recurrent_config=None,\n",
    "                          **kwargs):\n",
    "    net = nengo.Network(**kwargs)\n",
    "    if recurrent_config is None:\n",
    "        recurrent_config = nengo.Config(nengo.Connection)\n",
    "        recurrent_config[nengo.Connection].synapse = nengo.Lowpass(0.1)\n",
    "    with net:\n",
    "        net.ensemble = nengo.Ensemble(n_neurons, dimensions=dimensions + 1)\n",
    "        with recurrent_config:\n",
    "            nengo.Connection(\n",
    "                net.ensemble,\n",
    "                net.ensemble[:dimensions],\n",
    "                function=lambda x: x[:-1] * (1.0 - x[-1]))\n",
    "    return net\n",
    "\n",
    "\n",
    "def medial_pfc(coupling_strength,\n",
    "               n_neurons_per_integrator=200,\n",
    "               recurrent_config=None,\n",
    "               tau=0.1,\n",
    "               **kwargs):\n",
    "    net = nengo.Network(**kwargs)\n",
    "    with net:\n",
    "        recurrent_config = nengo.Config(nengo.Connection)\n",
    "        recurrent_config[nengo.Connection].synapse = nengo.Lowpass(tau)\n",
    "        net.pre = controlled_integrator(n_neurons_per_integrator, 1,\n",
    "                                        recurrent_config)\n",
    "        net.post = controlled_integrator(n_neurons_per_integrator, 1,\n",
    "                                         recurrent_config)\n",
    "        nengo.Connection(\n",
    "            net.pre.ensemble[0],\n",
    "            net.post.ensemble[0],\n",
    "            transform=coupling_strength)\n",
    "    return net\n",
    "\n",
    "\n",
    "def motor_cortex(command_threshold,\n",
    "                 n_neurons_per_command=30,\n",
    "                 ens_config=None,\n",
    "                 **kwargs):\n",
    "    net = nengo.Network(**kwargs)\n",
    "    if ens_config is None:\n",
    "        ens_config = nengo.Config(nengo.Ensemble)\n",
    "        ens_config[nengo.Ensemble].encoders = Choice([[1]])\n",
    "        ens_config[nengo.Ensemble].intercepts = Choice([command_threshold])\n",
    "    with net:\n",
    "        with ens_config:\n",
    "            net.press = nengo.Ensemble(n_neurons_per_command, dimensions=1)\n",
    "            net.release = nengo.Ensemble(n_neurons_per_command, dimensions=1)\n",
    "    return net\n",
    "\n",
    "\n",
    "def double_integrator(mpfc_coupling_strength,\n",
    "                      command_threshold,\n",
    "                      press_to_pre_gain=3,\n",
    "                      press_to_post_control=-6,\n",
    "                      recurrent_tau=0.1,\n",
    "                      **kwargs):\n",
    "    net = nengo.Network(**kwargs)\n",
    "    with net:\n",
    "        net.mpfc = medial_pfc(mpfc_coupling_strength)\n",
    "        net.motor = motor_cortex(command_threshold)\n",
    "        nengo.Connection(\n",
    "            net.motor.press,\n",
    "            net.mpfc.pre.ensemble[0],\n",
    "            transform=recurrent_tau * press_to_pre_gain)\n",
    "        nengo.Connection(\n",
    "            net.motor.press,\n",
    "            net.mpfc.post.ensemble[1],\n",
    "            transform=press_to_post_control)\n",
    "        nengo.Connection(net.mpfc.post.ensemble[0], net.motor.release)\n",
    "    return net\n",
    "\n",
    "\n",
    "def test_doubleintegrator(net):\n",
    "    # Provide input and probe outside of network construction,\n",
    "    # for more flexibility\n",
    "    with net:\n",
    "        nengo.Connection(\n",
    "            nengo.Node(lambda t: 1 if t < 0.2 else 0), net.motor.press)\n",
    "        pr_press = nengo.Probe(net.motor.press, synapse=0.01)\n",
    "        pr_release = nengo.Probe(net.motor.release, synapse=0.01)\n",
    "        pr_pre_int = nengo.Probe(net.mpfc.pre.ensemble[0], synapse=0.01)\n",
    "        pr_post_int = nengo.Probe(net.mpfc.post.ensemble[0], synapse=0.01)\n",
    "    with nengo.Simulator(net) as sim:\n",
    "        sim.run(1.4)\n",
    "    t = sim.trange()\n",
    "    plt.figure()\n",
    "    plt.subplot(2, 1, 1)\n",
    "    plt.plot(t, sim.data[pr_press], c='b', label=\"Press\")\n",
    "    plt.plot(t, sim.data[pr_release], c='g', label=\"Release\")\n",
    "    plt.axvspan(0, 0.2, color='b', alpha=0.3)\n",
    "    plt.axvspan(0.8, 1.2, color='g', alpha=0.3)\n",
    "    plt.xlim(right=1.4)\n",
    "    plt.legend(loc=\"best\")\n",
    "    plt.subplot(2, 1, 2)\n",
    "    plt.plot(t, sim.data[pr_pre_int], label=\"Pre Integrator\")\n",
    "    plt.plot(t, sim.data[pr_post_int], label=\"Post Integrator\")\n",
    "    plt.xlim(right=1.4)\n",
    "    plt.legend(loc=\"best\")\n",
    "\n",
    "\n",
    "for coupling_strength in (0.11, 0.16, 0.21):\n",
    "    # Try the same network with LIFRate neurons\n",
    "    with nengo.Config(nengo.Ensemble) as cfg:\n",
    "        cfg[nengo.Ensemble].neuron_type = nengo.LIFRate()\n",
    "        net = double_integrator(\n",
    "            mpfc_coupling_strength=coupling_strength,\n",
    "            command_threshold=0.85,\n",
    "            seed=0)\n",
    "    test_doubleintegrator(net)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
